import os
from random import sample
import torch
from EAattak.UniversalPerturbation import UAP
from EAattak.EvolutionAlgorithm import EA
from tqdm import tqdm, trange

class Attck:

    def __init__(self, dataer, neter, config):

        self.uaper = UAP(dataer=dataer, neter=neter, config=config)
        self.dataer = dataer
        self.neter = neter
        self.config = config
        self.root_path = './data/BaseSample'

    def target_attck(
            self, 
            origin_sample, 
            source_class, 
            target_class, 
            epsilon, 
            norm_type, 
            method, 
            batch, 
            size, 
            mode,
            limitation=10000,
        ):
        print('Origin Sample shape: {}, source_class {}, target_class {}'.format(origin_sample.shape, source_class, target_class))
        """
        origin_sample: 支持单张样本,且建议样本输入类型为tensor,带有batch=1的批次信息,例如(1, 3, 32, 32)
        """
        
        ## 检测样本字典是否存在
                
        # path = self.generate_path(epsilon, norm_type, method, batch, size, mode, isTarget=True)
        # path = os.path.join(
        #     self.root_path, 
        #     path,
        # )
        # if os.path.exists(path) == False:
        #     print('Base samples have not been generated at path <{}>'.format(path))
        #     print('Generating Base sample to path <{}>'.format(path))
        #     os.makedirs(path)
            
            ## this method is overwritten by the new version, 
            ## the original version can generate all target pairs, which is time-consuming
            # bs = self.uaper.generate_target_attck_universal_adversarial_perturbations(
            #     epsilon=epsilon,
            #     norm_type=norm_type,
            #     method=method,
            #     batch=batch,
            #     size=size,
            #     mode=mode,
            # )  ## TODO:: 完善UAP类, 如果是ImageNet类别的话，建议只生成想要的类别之间的样本
        
        self.file_check(
            epsilon=epsilon,
            norm_type=norm_type,
            method=method,
            batch=batch,
            size=size,
            mode=mode,
            isTarget=True,
        )

        path = self.generate_muaps_path(
            epsilon=epsilon,
            norm_type=norm_type,
            method=method,
            batch=batch,
            mode=mode,
            isTarget=True,
        )

        path = os.path.join(self.root_path, path)

        source_target_bs = {'perturbations' : []}

        ## 检测源类到目标类的映射是否存在
        source_target_file_name = '{}-{}'.format(source_class, target_class)

        if os.path.exists(os.path.join(path, source_target_file_name)) == False:
        
            print('The file do not have the source-target muaps, now generating...')
            os.makedirs(os.path.join(path, source_target_file_name))

            ## there is no muaps in the new files, now genrating from index: 0 to index: size - 1.
            for index in range(size):
                
                print('Generating the index {} MUAPs samples for ({})'.format(index, source_target_file_name))
                
                index_perturbation = self.uaper.generate_source_target_universal_adversarial_perturbations(
                    source_class=source_class,
                    target_class=target_class,
                    epsilon=epsilon,
                    norm_type=norm_type,
                    method=method,
                    batch=batch,
                    index=index, 
                    mode=mode,
                    size=size,
                )

                print('Saving the index {} uaps into path <{}>'.format(index, os.path.join(path, source_target_file_name, 'index_{}.pt'.format(index))))
                torch.save(index_perturbation, os.path.join(path, source_target_file_name, 'index_{}.pt'.format(index)))

                source_target_bs['perturbations'].append(index_perturbation)

        else:

            files = os.listdir(os.path.join(path, source_target_file_name))
            current_size = len(files)
            
            ## meets the size
            if current_size >= size:
                
                for index, file in enumerate(files):
                    
                    if index == size:
                        break

                    index_perturbation = torch.load(os.path.join(path, source_target_file_name, file))
                    source_target_bs['perturbations'].append(index_perturbation)

            ## current size leq size
            else:
                
                print('The storing uaps size is leq the goal size, now appending...')

                for file in files:

                    index_perturbation = torch.load(os.path.join(path, source_target_file_name, file))
                    source_target_bs['perturbations'].append(index_perturbation)

                for index in range(current_size, size):

                    print('Generating the index {} MUAPs samples for ({})'.format(index, source_target_file_name))
                
                    index_perturbation = self.uaper.generate_source_target_universal_adversarial_perturbations(
                        source_class=source_class,
                        target_class=target_class,
                        epsilon=epsilon,
                        norm_type=norm_type,
                        method=method,
                        batch=batch,
                        index=index, 
                        mode=mode,
                        size=size,
                    )

                    print('Saving the index {} uaps into path <{}>'.format(index, os.path.join(path, source_target_file_name, 'index_{}.pt'.format(index))))
                    torch.save(index_perturbation, os.path.join(path, source_target_file_name, 'index_{}.pt'.format(index)))

                    source_target_bs['perturbations'].append(index_perturbation)

        ## check the size
        if len(source_target_bs['perturbations']) != size:
            raise ValueError('The size of bs is not meet size, current: {}, Goal: {}.'.format(len(source_target_bs['perturbations'], size)))


        # if os.path.exists(os.path.join(path, source_target_file_name)) == False:
        #     print('The file do not have the source-target BaseSample, now generating...')
        #     source_target_bs = self.uaper.generate_source_target_universal_adversarial_perturbations(
        #         source_class=source_class,
        #         target_class=target_class,
        #         epsilon=epsilon,
        #         norm_type=norm_type,
        #         method=method,
        #         batch=batch,
        #         size=size,
        #         mode=mode,
        #     )
        # else:
        #     print('Load BaseSample from path <{}>'.format(os.path.join(path, source_target_file_name)))
        #     source_target_bs = torch.load(os.path.join(path, source_target_file_name))

        # ## 进入正式攻击阶段
        # ## 过度阶段，判断上述基样本是否已经可以攻击成功，并给攻击成功的样本提升优先级
        isSucc = False
        Succ_perturbation = None

        for index, perturbation in enumerate(source_target_bs['perturbations']):
            
            flag, index_list = self.neter.isTargetAttckSuccessful(origin_sample, perturbation, target_class)
            self.neter.query += 1
            if flag:
                isSucc = True
                Succ_perturbation = perturbation
   
        if isSucc == False:
            print('Stage I can not successfully attck the sample')
        ## 进化算法, stage II, 若过度阶段失败
        if isSucc == False:
            
            ## 定义评估函数. 然后喂入进化算法当中
            ## TODO, 接口仅以人口为输入
            ## 这里的评估函数，以最大化想扰动的样本的目标类概率为优化目标，进行构造目标函数
            def evaluation(populations):
                """
                给出每一个人口一个适应值，然后以向量的形式给出
                """
                ## 将种群转化生成一组扰动(映射)
                perturbations = self.dataer.transfer(populations)
                ##
                fitness = []

                for perturbation in perturbations:
                    adv_sample = torch.clamp(origin_sample + perturbation, min=0.0, max=1.0)
                    prediction = self.neter.get_prediction_score(adv_sample)
                    fitness.append(prediction.mean(dim=0)[target_class])
                return fitness


            optimizer = EA(
                dimesion=self.dataer.get_dim(),
                epsilon=epsilon,
                norm_type=norm_type,
                method=method,
                evaluation=evaluation,
                config=self.config,
                optimizer='GA',
                strategy='base',
                population_size=self.config['population_size'],
                perturbations=[self.dataer.retransfer(perturbation) for perturbation in source_target_bs['perturbations']],
                isExtension=True,
            )
            query_count = 0
            query_count += optimizer.population_size
            
            with trange(self.config['generation']) as t:
                for gen in t:
                    t.set_description('GEN {}'.format(gen))
                    if isSucc:
                        break
                    if query_count > limitation:
                        break
                    optimizer.step()
                    query_count += optimizer.population_size
                    
                    ## TODO bottleneck problem
                    maxValue = max(optimizer.fitness)
                    t.set_postfix(probability='{:.2f}%'.format(maxValue * 100))
                     
                    perturbations = self.dataer.transfer(optimizer.populations)
                    flag, index_list = self.neter.isTargetAttckSuccessful(origin_sample, perturbations, target_class)
                    if flag:
                        isSucc = True
                        Succ_perturbation = perturbations[index_list[0]]  ## TODO 这里就给一个样本的成功扰动就行
                        break
            
        ## Finally
        del source_target_bs
        
        if isSucc:
            return isSucc, torch.clamp(origin_sample + Succ_perturbation, min=0.0, max=1.0)
        else:
            return isSucc, None

    def untarget_attack(
            self, 
            origin_sample, 
            source_class, 
            epsilon, 
            norm_type, 
            method, 
            batch, 
            size, 
            mode,
            limitation=10000,
        ):
        print('Origin Sample shape: {}, source_class {}'.format(origin_sample.shape, source_class))
        """
        origin_sample: 支持单张样本，且建议样本输入类型为tensor，带有batch=1的批次信息，例如(1, 3, 32, 32)
        """
        
        ## 检测样本字典是否存在
        # path = self.generate_path(epsilon, norm_type, method, batch, size, mode, isTarget=False)
        # path = os.path.join(
        #     self.root_path, 
        #     path,
        # )
        # if os.path.exists(path) == False:
        #     print('Base samples have not been generated at path <{}>'.format(path))
        #     print('Generating Base sample to path <{}>'.format(path))
        #     os.makedirs(path)
            
        # ## 检测源类映射是否存在
        # source_file_name = '{}.pt'.format(source_class)
        # if os.path.exists(os.path.join(path, source_file_name)) == False:
        #     print('The file do not have the source BaseSample, now generating...')
        #     source_bs = self.uaper.generate_untarget_universal_adversarial_perturbations(
        #         source_class=source_class,
        #         epsilon=epsilon,
        #         norm_type=norm_type,
        #         method=method,
        #         batch=batch,
        #         size=size,
        #         mode=mode,
        #     )
        # else:
        #     print('Load BaseSample from path <{}>'.format(os.path.join(path, source_file_name)))
        #     source_bs = torch.load(os.path.join(path, source_file_name))


        
        self.file_check(
            epsilon=epsilon,
            norm_type=norm_type,
            method=method,
            batch=batch,
            size=size,
            mode=mode,
            isTarget=False,
        )

        path = self.generate_muaps_path(
            epsilon=epsilon,
            norm_type=norm_type,
            method=method,
            batch=batch,
            mode=mode,
            isTarget=False,
        )

        path = os.path.join(self.root_path, path)

        source_bs = {'perturbations' : []}

        ## 检测源类到目标类的映射是否存在
        source_file_name = '{}'.format(source_class)

        if os.path.exists(os.path.join(path, source_file_name)) == False:
        
            print('The file do not have the source muaps, now generating...')
            os.makedirs(os.path.join(path, source_file_name))

            ## there is no muaps in the new files, now genrating from index: 0 to index: size - 1.
            for index in range(size):
                
                print('Generating the index {} MUAPs samples for ({})'.format(index, source_file_name))
                
                index_perturbation = self.uaper.generate_untarget_universal_adversarial_perturbations(
                    source_class=source_class,
                    epsilon=epsilon,
                    norm_type=norm_type,
                    method=method,
                    batch=batch,
                    index=index,
                    mode=mode,
                    size=size, 
                )

                print('Saving the index {} uaps into path <{}>'.format(index, os.path.join(path, source_file_name, 'index_{}.pt'.format(index))))
                torch.save(index_perturbation, os.path.join(path, source_file_name, 'index_{}.pt'.format(index)))

                source_bs['perturbations'].append(index_perturbation)

        else:

            files = os.listdir(os.path.join(path, source_file_name))
            current_size = len(files)
            
            ## meets the size
            if current_size >= size:
                
                for index, file in enumerate(files):
                    
                    if index == size:
                        break

                    index_perturbation = torch.load(os.path.join(path, source_file_name, file))
                    source_bs['perturbations'].append(index_perturbation)

            ## current size leq size
            else:
                
                print('The storing uaps size is leq the goal size, now appending...')

                for file in files:

                    index_perturbation = torch.load(os.path.join(path, source_file_name, file))
                    source_bs['perturbations'].append(index_perturbation)

                for index in range(current_size, size):

                    print('Generating the index {} MUAPs samples for ({})'.format(index, source_file_name))
                
                    index_perturbation = self.uaper.generate_untarget_universal_adversarial_perturbations(
                        source_class=source_class,
                        epsilon=epsilon,
                        norm_type=norm_type,
                        method=method,
                        batch=batch,
                        index=index,
                        mode=mode,
                        size=size, 
                    )

                    print('Saving the index {} uaps into path <{}>'.format(index, os.path.join(path, source_file_name, 'index_{}.pt'.format(index))))
                    torch.save(index_perturbation, os.path.join(path, source_file_name, 'index_{}.pt'.format(index)))

                    source_bs['perturbations'].append(index_perturbation)

        ## check the size
        if len(source_bs['perturbations']) != size:
            raise ValueError('The size of bs is not meet size, current: {}, Goal: {}.'.format(len(source_bs['perturbations'], size)))

        # ## 进入正式攻击阶段
        # ## 过度阶段，判断上述基样本是否已经可以攻击成功，并给攻击成功的样本提升优先级
        isSucc = False
        Succ_perturbation = None

        for index, perturbation in enumerate(source_bs['perturbations']):
            
            flag, index_list = self.neter.isUntargetAttackSuccessful(origin_sample, perturbation, source_class)
            self.neter.query += 1
            if flag:
                isSucc = True
                Succ_perturbation = perturbation

        if isSucc == False:
            print('Stage I can not successfully attck the sample')
        ## 进化算法, stage II, 若过度阶段失败
        if isSucc == False:
            
            ## 定义评估函数. 然后喂入进化算法当中
            ## TODO, 接口仅以人口为输入
            ## 这里的评估函数，以最大化想扰动的样本的目标类概率为优化目标，进行构造目标函数
            def evaluation(populations):
                """
                给出每一个人口一个适应值，然后以向量的形式给出
                """
                ## 将种群转化生成一组扰动(映射)
                perturbations = self.dataer.transfer(populations)
                ##
                fitness = []

                for perturbation in perturbations:
                    adv_sample = torch.clamp(origin_sample + perturbation, min=0.0, max=1.0)
                    prediction = self.neter.get_prediction_score(adv_sample)
                    fitness.append(-prediction.mean(dim=0)[source_class])
                return fitness

            optimizer = EA(
                dimesion=self.dataer.get_dim(),
                epsilon=epsilon,
                norm_type=norm_type,
                method=method,
                evaluation=evaluation,
                config=self.config,
                optimizer='GA',
                strategy='base',
                population_size=self.config['population_size'],
                perturbations=[self.dataer.retransfer(perturbation) for perturbation in source_bs['perturbations']],
                isExtension=True,
            )
            query_count = 0
            query_count += optimizer.population_size

            with trange(self.config['generation']) as t:
                for gen in t:
                    t.set_description('GEN {}'.format(gen))
                    if isSucc:
                        break
                    if query_count > limitation:
                        break
                    optimizer.step()
                    query_count += optimizer.population_size

                    ## TODO bottleneck problem
                    maxValue = max(optimizer.fitness)
                    t.set_postfix(probability='{:.2f}%'.format(-1 * maxValue * 100))
                     
                    perturbations = self.dataer.transfer(optimizer.populations)
                    flag, index_list = self.neter.isUntargetAttackSuccessful(origin_sample, perturbations, source_class)
                    if flag:
                        isSucc = True
                        Succ_perturbation = perturbations[index_list[0]]  ## TODO 这里就给一个样本的成功扰动就行
                        break
            
        ## Finally
        del source_bs
        
        if isSucc:
            return isSucc, torch.clamp(origin_sample + Succ_perturbation, min=0.0, max=1.0)
        else:
            return isSucc, None

    def untarget_MPack_attack(
            self,
            origin_sample, 
            source_class, 
            epsilon, 
            norm_type, 
            method, 
            batch, 
            size, 
            mode
        ):
        print('Origin Sample shape: {}, source_class {}'.format(origin_sample.shape, source_class))
        """
        origin_sample: 支持单张样本，且建议样本输入类型为tensor，带有batch=1的批次信息，例如(1, 3, 32, 32)
        """
        
        ## 检测样本字典是否存在
        # path = self.generate_path(epsilon, norm_type, method, batch, size, mode, isTarget=False)
        # path = os.path.join(
        #     self.root_path, 
        #     path,
        # )
        # if os.path.exists(path) == False:
        #     print('Base samples have not been generated at path <{}>'.format(path))
        #     print('Generating Base sample to path <{}>'.format(path))
        #     os.makedirs(path)
            
        # ## 检测源类映射是否存在
        # source_file_name = '{}.pt'.format(source_class)
        # if os.path.exists(os.path.join(path, source_file_name)) == False:
        #     print('The file do not have the source BaseSample, now generating...')
        #     source_bs = self.uaper.generate_untarget_universal_adversarial_perturbations(
        #         source_class=source_class,
        #         epsilon=epsilon,
        #         norm_type=norm_type,
        #         method=method,
        #         batch=batch,
        #         size=size,
        #         mode=mode,
        #     )
        # else:
        #     print('Load BaseSample from path <{}>'.format(os.path.join(path, source_file_name)))
        #     source_bs = torch.load(os.path.join(path, source_file_name))

        
        self.file_check(
            epsilon=epsilon,
            norm_type=norm_type,
            method=method,
            batch=batch,
            size=size,
            mode=mode,
            isTarget=False,
        )

        path = self.generate_muaps_path(
            epsilon=epsilon,
            norm_type=norm_type,
            method=method,
            batch=batch,
            mode=mode,
            isTarget=False,
        )

        path = os.path.join(self.root_path, path)

        source_bs = {'perturbations' : []}

        ## 检测源类到目标类的映射是否存在
        source_file_name = '{}'.format(source_class)

        if os.path.exists(os.path.join(path, source_file_name)) == False:
        
            print('The file do not have the source muaps, now generating...')
            os.makedirs(os.path.join(path, source_file_name))

            ## there is no muaps in the new files, now genrating from index: 0 to index: size - 1.
            for index in range(size):
                
                print('Generating the index {} MUAPs samples for ({})'.format(index, source_file_name))
                
                index_perturbation = self.uaper.generate_untarget_universal_adversarial_perturbations(
                    source_class=source_class,
                    epsilon=epsilon,
                    norm_type=norm_type,
                    method=method,
                    batch=batch,
                    index=index,
                    mode=mode,
                    size=size, 
                )

                print('Saving the index {} uaps into path <{}>'.format(index, os.path.join(path, source_file_name, 'index_{}.pt'.format(index))))
                torch.save(index_perturbation, os.path.join(path, source_file_name, 'index_{}.pt'.format(index)))

                source_bs['perturbations'].append(index_perturbation)

        else:

            files = os.listdir(os.path.join(path, source_file_name))
            current_size = len(files)

            ## meets the size
            if current_size >= size:
                
                for index, file in enumerate(files):
                    
                    if index == size:
                        break

                    index_perturbation = torch.load(os.path.join(path, source_file_name, file))
                    source_bs['perturbations'].append(index_perturbation)

            ## current size leq size
            else:
                
                print('The storing uaps size is leq the goal size, now appending...')

                for file in files:

                    index_perturbation = torch.load(os.path.join(path, source_file_name, file))
                    source_bs['perturbations'].append(index_perturbation)

                for index in range(current_size, size):

                    print('Generating the index {} MUAPs samples for ({})'.format(index, source_file_name))
                
                    index_perturbation = self.uaper.generate_untarget_universal_adversarial_perturbations(
                        source_class=source_class,
                        epsilon=epsilon,
                        norm_type=norm_type,
                        method=method,
                        batch=batch,
                        index=index,
                        mode=mode,
                        size=size, 
                    )

                    print('Saving the index {} uaps into path <{}>'.format(index, os.path.join(path, source_file_name, 'index_{}.pt'.format(index))))
                    torch.save(index_perturbation, os.path.join(path, source_file_name, 'index_{}.pt'.format(index)))

                    source_bs['perturbations'].append(index_perturbation)

        ## check the size
        if len(source_bs['perturbations']) != size:
            raise ValueError('The size of bs is not meet size, current: {}, Goal: {}.'.format(len(source_bs['perturbations'], size)))

        # ## 进入正式攻击阶段
        # ## 过度阶段，判断上述基样本是否已经可以攻击成功，并给攻击成功的样本提升优先级
        isSucc = False
        Succ_perturbation = None

        for index, perturbation in enumerate(source_bs['perturbations']):
            
            flag, index_list = self.neter.isUntargetAttackSuccessful(origin_sample, perturbation, source_class)
            self.neter.query += 1
            if flag:
                isSucc = True
                Succ_perturbation = perturbation

        ## Finally
        del source_bs
        
        if isSucc:
            return isSucc, torch.clamp(origin_sample + Succ_perturbation, min=0.0, max=1.0)
        else:
            return isSucc, None

    def target_MPack_attack(
            self, 
            origin_sample, 
            source_class, 
            target_class, 
            epsilon, 
            norm_type, 
            method, 
            batch, 
            size, 
            mode
        ):
        print('Origin Sample shape: {}, source_class {}, target_class {}'.format(origin_sample.shape, source_class, target_class))
        """
        origin_sample: 支持单张样本，且建议样本输入类型为tensor，带有batch=1的批次信息，例如(1, 3, 32, 32)
        """
        
        ## 检测样本字典是否存在
        # path = self.generate_path(epsilon, norm_type, method, batch, size, mode, isTarget=True)
        # path = os.path.join(
        #     self.root_path, 
        #     path,
        # )
        # if os.path.exists(path) == False:
        #     print('Base samples have not been generated at path <{}>'.format(path))
        #     print('Generating Base sample to path <{}>'.format(path))
        #     os.makedirs(path)
            
        #     ## this method is overwritten by the new version, 
        #     ## the original version can generate all target pairs, which is time-consuming
        #     # bs = self.uaper.generate_target_attck_universal_adversarial_perturbations(
        #     #     epsilon=epsilon,
        #     #     norm_type=norm_type,
        #     #     method=method,
        #     #     batch=batch,
        #     #     size=size,
        #     #     mode=mode,
        #     # )  ## TODO:: 完善UAP类, 如果是ImageNet类别的话，建议只生成想要的类别之间的样本

        # ## 检测源类到目标类的映射是否存在
        # source_target_file_name = '{}-{}.pt'.format(source_class, target_class)
        # if os.path.exists(os.path.join(path, source_target_file_name)) == False:
        #     print('The file do not have the source-target BaseSample, now generating...')
        #     source_target_bs = self.uaper.generate_source_target_universal_adversarial_perturbations(
        #         source_class=source_class,
        #         target_class=target_class,
        #         epsilon=epsilon,
        #         norm_type=norm_type,
        #         method=method,
        #         batch=batch,
        #         size=size,
        #         mode=mode,
        #     )
        # else:
        #     print('Load BaseSample from path <{}>'.format(os.path.join(path, source_target_file_name)))
        #     source_target_bs = torch.load(os.path.join(path, source_target_file_name))

        self.file_check(
            epsilon=epsilon,
            norm_type=norm_type,
            method=method,
            batch=batch,
            size=size,
            mode=mode,
            isTarget=True,
        )

        path = self.generate_muaps_path(
            epsilon=epsilon,
            norm_type=norm_type,
            method=method,
            batch=batch,
            mode=mode,
            isTarget=True,
        )

        path = os.path.join(self.root_path, path)

        source_target_bs = {'perturbations' : []}

        ## 检测源类到目标类的映射是否存在
        source_target_file_name = '{}-{}'.format(source_class, target_class)

        if os.path.exists(os.path.join(path, source_target_file_name)) == False:
        
            print('The file do not have the source-target muaps, now generating...')
            os.makedirs(os.path.join(path, source_target_file_name))

            ## there is no muaps in the new files, now genrating from index: 0 to index: size - 1.
            for index in range(size):
                
                print('Generating the index {} MUAPs samples for ({})'.format(index, source_target_file_name))
                
                index_perturbation = self.uaper.generate_source_target_universal_adversarial_perturbations(
                    source_class=source_class,
                    target_class=target_class,
                    epsilon=epsilon,
                    norm_type=norm_type,
                    method=method,
                    batch=batch,
                    index=index, 
                    mode=mode,
                    size=size,
                )

                print('Saving the index {} uaps into path <{}>'.format(index, os.path.join(path, source_target_file_name, 'index_{}.pt'.format(index))))
                torch.save(index_perturbation, os.path.join(path, source_target_file_name, 'index_{}.pt'.format(index)))

                source_target_bs['perturbations'].append(index_perturbation)

        else:

            files = os.listdir(os.path.join(path, source_target_file_name))
            current_size = len(files)
            
            ## meets the size
            if current_size >= size:
                
                for index, file in enumerate(files):
                    
                    if index == size:
                        break

                    index_perturbation = torch.load(os.path.join(path, source_target_file_name, file))
                    source_target_bs['perturbations'].append(index_perturbation)

            ## current size leq size
            else:
                
                print('The storing uaps size is leq the goal size, now appending...')

                for file in files:

                    index_perturbation = torch.load(os.path.join(path, source_target_file_name, file))
                    source_target_bs['perturbations'].append(index_perturbation)

                for index in range(current_size, size):

                    print('Generating the index {} MUAPs samples for ({})'.format(index, source_target_file_name))
                
                    index_perturbation = self.uaper.generate_source_target_universal_adversarial_perturbations(
                        source_class=source_class,
                        target_class=target_class,
                        epsilon=epsilon,
                        norm_type=norm_type,
                        method=method,
                        batch=batch,
                        index=index, 
                        mode=mode,
                        size=size,
                    )

                    print('Saving the index {} uaps into path <{}>'.format(index, os.path.join(path, source_target_file_name, 'index_{}.pt'.format(index))))
                    torch.save(index_perturbation, os.path.join(path, source_target_file_name, 'index_{}.pt'.format(index)))

                    source_target_bs['perturbations'].append(index_perturbation)

        ## check the size
        if len(source_target_bs['perturbations']) != size:
            raise ValueError('The size of bs is not meet size, current: {}, Goal: {}.'.format(len(source_target_bs['perturbations'], size)))

        # ## 进入正式攻击阶段
        # ## 过度阶段，判断上述基样本是否已经可以攻击成功，并给攻击成功的样本提升优先级
        isSucc = False
        Succ_perturbation = None

        for index, perturbation in enumerate(source_target_bs['perturbations']):
            
            flag, index_list = self.neter.isTargetAttckSuccessful(origin_sample, perturbation, target_class)
            self.neter.query += 1
            if flag:
                isSucc = True
                Succ_perturbation = perturbation

        ## Finally
        del source_target_bs
        
        if isSucc:
            return isSucc, torch.clamp(origin_sample + Succ_perturbation, min=0.0, max=1.0)
        else:
            return isSucc, None

    def untarget_GA_attack(
            self, 
            origin_sample, 
            source_class, 
            epsilon, 
            norm_type, 
            method, 
            limitation=10000,
        ):
        print('Origin Sample shape: {}, source_class {}'.format(origin_sample.shape, source_class))
        """
        origin_sample: 支持单张样本，且建议样本输入类型为tensor，带有batch=1的批次信息，例如(1, 3, 32, 32)
        """
  
        # ## 进入正式攻击阶段
        # ## 过度阶段，判断上述基样本是否已经可以攻击成功，并给攻击成功的样本提升优先级
        query_count = 0
        isSucc = False
        Succ_perturbation = None
        
        ## 定义评估函数. 然后喂入进化算法当中
        ## TODO, 接口仅以人口为输入
        ## 这里的评估函数，以最大化想扰动的样本的目标类概率为优化目标，进行构造目标函数
        def evaluation(populations):
            """
            给出每一个人口一个适应值，然后以向量的形式给出
            """
            ## 将种群转化生成一组扰动(映射)
            perturbations = self.dataer.transfer(populations)
            ##
            fitness = []

            for perturbation in perturbations:
                adv_sample = torch.clamp(origin_sample + perturbation, min=0.0, max=1.0)
                prediction = self.neter.get_prediction_score(adv_sample)
                fitness.append(-prediction.mean(dim=0)[source_class])
            return fitness

        optimizer = EA(
            dimesion=self.dataer.get_dim(),
            epsilon=epsilon,
            norm_type=norm_type,
            method=method,
            evaluation=evaluation,
            config=self.config,
            optimizer='GA',
            strategy='random',
            population_size=self.config['population_size'],
        )

        query_count += optimizer.population_size
        
        with trange(self.config['generation']) as t:
            for gen in t:
                t.set_description('GEN {}'.format(gen))
                if isSucc:
                    break
                if query_count > limitation:
                    break
                optimizer.step()

                query_count += optimizer.population_size
                
                ## TODO bottleneck problem
                maxValue = max(optimizer.fitness)
                t.set_postfix(probability='{:.2f}%'.format(-1 * maxValue * 100))
                    
                perturbations = self.dataer.transfer(optimizer.populations)
                flag, index_list = self.neter.isUntargetAttackSuccessful(origin_sample, perturbations, source_class)
                if flag:
                    isSucc = True
                    Succ_perturbation = perturbations[index_list[0]]  ## TODO 这里就给一个样本的成功扰动就行
                    break
        
        if isSucc:
            return isSucc, torch.clamp(origin_sample + Succ_perturbation, min=0.0, max=1.0)
        else:
            return isSucc, None
        
    def target_GA_attack( 
            self, 
            origin_sample, 
            source_class, 
            target_class, 
            epsilon, 
            norm_type, 
            method, 
            limitation=10000,
        ):
        print('Origin Sample shape: {}, source_class {}, target_class {}'.format(origin_sample.shape, source_class, target_class))
        """
        origin_sample: 支持单张样本，且建议样本输入类型为tensor，带有batch=1的批次信息，例如(1, 3, 32, 32)
        """
        query_count = 0
        isSucc = False
        Succ_perturbation = None
        
        ## 定义评估函数. 然后喂入进化算法当中
        ## TODO, 接口仅以人口为输入
        ## 这里的评估函数，以最大化想扰动的样本的目标类概率为优化目标，进行构造目标函数
        def evaluation(populations):
            """
            给出每一个人口一个适应值，然后以向量的形式给出
            """
            ## 将种群转化生成一组扰动(映射)
            perturbations = self.dataer.transfer(populations)
            ##
            fitness = []

            for perturbation in perturbations:
                adv_sample = torch.clamp(origin_sample + perturbation, min=0.0, max=1.0)
                prediction = self.neter.get_prediction_score(adv_sample)
                fitness.append(prediction.mean(dim=0)[target_class])
            return fitness

        optimizer = EA(
            dimesion=self.dataer.get_dim(),
            epsilon=epsilon,
            norm_type=norm_type,
            method=method,
            evaluation=evaluation,
            config=self.config,
            optimizer='GA',
            strategy='random',
            population_size=self.config['population_size'],
        )

        query_count += optimizer.population_size

        with trange(self.config['generation']) as t:
            for gen in t:
                t.set_description('GEN {}'.format(gen))
                if isSucc:
                    break
                if query_count > limitation:
                    break

                optimizer.step()
                query_count += optimizer.population_size
                
                ## TODO bottleneck problem
                maxValue = max(optimizer.fitness)
                t.set_postfix(probability='{:.2f}%'.format(maxValue * 100))
                    
                perturbations = self.dataer.transfer(optimizer.populations)
                flag, index_list = self.neter.isTargetAttckSuccessful(origin_sample, perturbations, target_class)
                if flag:
                    isSucc = True
                    Succ_perturbation = perturbations[index_list[0]]  ## TODO 这里就给一个样本的成功扰动就行
                    break
        
        if isSucc:
            return isSucc, torch.clamp(origin_sample + Succ_perturbation, min=0.0, max=1.0)
        else:
            return isSucc, None

    def generate_path(
            self, 
            epsilon, 
            norm_type, 
            method, 
            batch, 
            size, 
            mode, 
            isTarget=True
        ):
        """
        return file path and father path
        """
        if isTarget:
            father_path = '{}/{}/Target/Batch({})_Size({})_Mode({})'.format(self.config['model_arch'], self.config['dataset'], batch, size, mode)
            path = os.path.join(father_path, 'Eps({:.4f})_Norm({})_Method({})'.format(epsilon, norm_type, method))
            return path
        else:
            father_path = '{}/{}/Untarget/Batch({})_Size({})_Mode({})'.format(self.config['model_arch'], self.config['dataset'], batch, size, mode)
            path = os.path.join(father_path, 'Eps({:.4f})_Norm({})_Method({})'.format(epsilon, norm_type, method))
            return path
    
    def generate_muaps_path(
        self,
        epsilon,
        norm_type,
        method,
        batch, 
        mode,
        isTarget=True
        ):
        
        if isTarget:
            father_path = '{}/{}/Target/Batch({})_Mode({})'.format(self.config['model_arch'], self.config['dataset'], batch, mode)
            path = os.path.join(father_path, 'Eps({:.4f})_Norm({})_Method({})'.format(epsilon, norm_type, method))
            return path
        else:
            father_path = '{}/{}/Untarget/Batch({})_Mode({})'.format(self.config['model_arch'], self.config['dataset'], batch, mode)
            path = os.path.join(father_path, 'Eps({:.4f})_Norm({})_Method({})'.format(epsilon, norm_type, method))
            return path
    
    ## Need to rewrite
    def untarget_UAP_show(
            self,
            origin_sample, 
            source_class, 
            epsilon, 
            norm_type, 
            method, 
            batch, 
            size, 
            mode
        ):
        print('Origin Sample shape: {}, source_class {}'.format(origin_sample.shape, source_class))
        """
        origin_sample: 支持单张样本，且建议样本输入类型为tensor，带有batch=1的批次信息，例如(1, 3, 32, 32)
        """
        
        ## 检测样本字典是否存在
        path = self.generate_path(epsilon, norm_type, method, batch, size, mode, isTarget=False)
        path = os.path.join(
            self.root_path, 
            path,
        )
        if os.path.exists(path) == False:
            print('Base samples have not been generated at path <{}>'.format(path))
            print('Generating Base sample to path <{}>'.format(path))
            os.makedirs(path)
            
        ## 检测源类映射是否存在
        source_file_name = '{}.pt'.format(source_class)
        if os.path.exists(os.path.join(path, source_file_name)) == False:
            print('The file do not have the source BaseSample, now generating...')
            source_bs = self.uaper.generate_untarget_universal_adversarial_perturbations(
                source_class=source_class,
                epsilon=epsilon,
                norm_type=norm_type,
                method=method,
                batch=batch,
                size=size,
                mode=mode,
            )
        else:
            print('Load BaseSample from path <{}>'.format(os.path.join(path, source_file_name)))
            source_bs = torch.load(os.path.join(path, source_file_name))

        Succ_list = []

        for index, perturbation in enumerate(source_bs['perturbations']):
            
            flag, index_list = self.neter.isUntargetAttackSuccessful(origin_sample, perturbation, source_class)
            self.neter.query += 1
            if flag:
                Succ_list.append(index)

        ## Finally
        del source_bs
        
        return Succ_list

    # Need to rewrite
    def get_untarget_MUAPs(
            self,
            source_class, 
            epsilon, 
            norm_type, 
            method, 
            batch, 
            size, 
            mode
        ):
        
        ## 检测样本字典是否存在
        path = self.generate_path(epsilon, norm_type, method, batch, size, mode, isTarget=False)
        path = os.path.join(
            self.root_path, 
            path,
        )
        if os.path.exists(path) == False:
            print('Base samples have not been generated at path <{}>'.format(path))
            print('Generating Base sample to path <{}>'.format(path))
            os.makedirs(path)
            
        ## 检测源类映射是否存在
        source_file_name = '{}.pt'.format(source_class)
        if os.path.exists(os.path.join(path, source_file_name)) == False:
            print('The file do not have the source BaseSample, now generating...')
            source_bs = self.uaper.generate_untarget_universal_adversarial_perturbations(
                source_class=source_class,
                epsilon=epsilon,
                norm_type=norm_type,
                method=method,
                batch=batch,
                size=size,
                mode=mode,
            )
        else:
            print('Load BaseSample from path <{}>'.format(os.path.join(path, source_file_name)))
            source_bs = torch.load(os.path.join(path, source_file_name))

        perturbations = torch.cat(source_bs['perturbations'], dim=0)
     
        return perturbations
    
    ## Undo
    def batch_target_attck(self, ):
        """
        支持批量的目标攻击，后续写
        """
        pass

    ## Undo
    def batch_untarget_attck(self, ):

        """
        支持 批量的非目标攻击，后续写
        """
        pass

    def file_check(
        self,
        epsilon, 
        norm_type, 
        method, 
        batch, 
        size, 
        mode, 
        isTarget, 
        ):

        ## 检测新老路径
        ## case0: 新路径存在，直接pass，用新路径
        ## case1: 老路径不存在，新路径不存在。创建直接创建新路径
        ## case3: 老路径存在，新路径不存在。 创建新路径，将老路径的数据转化为新路径的数据 

        new_path = os.path.join(
            self.root_path, 
            self.generate_muaps_path(
                epsilon=epsilon,
                norm_type=norm_type,
                method=method,
                batch=batch,
                mode=mode,
                isTarget=isTarget,
            )
        )
        if os.path.exists(new_path) == True:
            
            print('Find the new type MUAPs files')
            return True
        
        else:
            print('Not find new tupe MUAPs files, now create the file path...')
            os.makedirs(new_path)


        old_path = os.path.join(
            self.root_path, 
            self.generate_path(
                epsilon=epsilon,
                norm_type=norm_type,
                method=method,
                batch=batch,
                size=10,  ## old way usually uses size=16 to generate, use 16 to laod and transfer old data.
                mode=mode,
                isTarget=isTarget,
            )
        )

        if os.path.exists(old_path) == False:

            print('Not find old type MUAPs files')
            return False
        
        else:

            print('Now tranfer the old MUAPs into new type files...')

            print('The old path<{}>'.format(old_path))

            for root, dirs, files in os.walk(old_path):
                for filename in files:
                    ## print(filename[:-3])
                    print('Loading old file <{}>'.format(os.path.join(root, filename)))
                    print('Creating new file <{}>'.format(os.path.join(new_path, filename[:-3])))
                    os.makedirs(os.path.join(new_path, filename[:-3]))
                    print('Create done ! Now loading data...')

                    dicter = torch.load(os.path.join(old_path, filename))

                    for index, perturbation in enumerate(dicter['perturbations']):
                        
                        path = os.path.join(new_path, filename[:-3], 'index_{}.pt'.format(index))
                        print('Saving into path <{}>'.format(path))
                        torch.save(perturbation, path)

                    print('Save done !')
            
            print('Transfering from old path into new path has done !')



        